rm(list=ls())
gc()
sample_name <- "300s"
savefolder_base <- paste0(getwd(), "/figures", sample_name, "/")
fit_bayes_models <- FALSE  # fitted objects are in replication materials

### Load sources
source("functions.R")

stan_iteration <- 1500
analysis_common <- list(savefolder_base, sample_name)


### Import libraries
library(tidyverse)
library(gridExtra)
library(rstan)
library(ggmcmc)
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())


##
## Load Data
##
#read_raw_data("../CleanedFinaldata.csv") %>% # remove identifying information
#  write_csv("DistFinalData.csv")
raw_data <- read_csv("DistFinalData.csv")

if(sample_name == "300s"){
  data <- clean_data(raw_data, min_time = 300, max_time = 40000, time = "GroupQ")
} else if(sample_name == "dist"){

  time <- raw_data %>% filter(completed == 1) %>% 
            mutate(Time = Debriefing_Start_Time - KnowledgeQ_Start_Time) %>%
            filter(Time > 0) %>% 
            pull(Time) %>% sort()

  use_time <- time[as.integer(length(time)*0.10):as.integer(length(time)*0.95)]
  data <- clean_data(raw_data, min_time = as.integer(min(use_time)), 
                                max_time = as.integer(max(use_time)), time = "GroupQ")
            
}else{
  stop("Please specify sample type")
}


##
## Analysis (this section takes time, fitted objects are in replication materials)
##
if (fit_bayes_models) {
  ### Load Stan models
  # Without Focus
  ordered_logit <- stan_model(file = "ordered_logit.stan", model_name = "ologit1")
  ordered_logit_noCov <- stan_model(file = "ordered_logit_noCov.stan", model_name = "ologit1noCov")

  # With Focus
  ordered_logit2 <- stan_model(file = "ordered_logit2.stan", model_name = "ologit2")
  ordered_logit2_noCov <- stan_model(file = "ordered_logit2_noCov.stan", model_name = "ologit2noCov")

  # Main analysis
  main_analysis(data,
               choice_rights, 
               cov_types, cov_types_names,
               interaction_types, interaction_names,
               analysis_functions, figure_functions
             )
}


## Figures
hdi_interval <- 0.90

use_order <- c("A9DefenceAmy", "A9IntCoop", "A9TerrResor",
               "EmergencyDecree", "EmergencyDissolution", "EmergencyOrder",
               "Family", "Religion", "VotingRights", "Environment", "Privacy",
               "Transparency")

use_order_name <- c("Art 9: Defense army",
                    "Art 9: International cooperation",
                    "Art 9: Territory and resources",
                    "Emergency: Decree",
                    "Emergency: No HR dissolution",
                    "Emergency: Follow orders",
                    "Family duty",
                    "Separation of church and state",
                    "Foreigners' voting rights",
                    "Environment",
                    "Privacy",
                    "Government transparency")

foldername <- paste0(savefolder_base, "result_tile/")
if(!file.exists(foldername)){
  message("Figure folder does not exist. Created.")  
  dir.create(foldername)
}

plot_res(analysis = "_WithoutIncome_None_BayesCov")
figure_all_choices(analysis = "_WithoutIncome_None_BayesCov")

plot_res_focus(analysis = "_WithoutIncome_IdealConst_BayesCovFocus", focus_name = c("Protect Rights", "Embody Ideals"))
plot_res_focus(analysis = "_WithoutIncome_partyLDP_BayesCovFocus", focus_name = c("Non-LDP", "LDP"))
plot_res_focus(analysis = "_WithoutIncome_partyIndp_BayesCovFocus", focus_name = c("Non-Independent", "Independent"))
plot_res_focus(analysis = "_WithoutIncome_partyLeft_BayesCovFocus", focus_name = c("Non-Left", "Left"))
plot_res_focus(analysis = "_WithoutIncome_Knowledge_BayesCovFocus", focus_name = c("Unknowledgeable", "Knowledgeable"))


## Create Table 2
table_res("diff_D")
table_res("diff_D", significant_only = FALSE)


##
## Descriptive Statistics
##
tmpdata <- data_description(data)


## Knowledge questions
table(tmpdata$Q3_CorrectNum)

## RHat
check_rhat()


## Footnote 18
broom::tidy(lm(Party_LDP ~ Group, data = tmpdata)) %>% filter(term == "GroupB") %>% pull(p.value)
broom::tidy(lm(Party_Left ~ Group, data = tmpdata)) %>% filter(term == "GroupB") %>% pull(p.value)
broom::tidy(lm(Party_Indp ~ Group, data = tmpdata)) %>% filter(term == "GroupB") %>% pull(p.value)
broom::tidy(lm(Q14.3 ~ Group, data = tmpdata %>% mutate(Q14.3 = if_else(Q14.3 == "Embody Ideals", 1, 0)))) %>% 
  filter(term == "GroupB") %>% pull(p.value)


## Idealists / pragmatists
table(tmpdata$PartyID, tmpdata$Q14.3) %>% knitr::kable()
table(tmpdata$Q3_CorrectNum, tmpdata$Q14.3) %>% knitr::kable()


## Party ID and Constitutional Knowledge
broom::tidy(lm(Q3_CorrectNum ~ Party_LDP, data = tmpdata))
broom::tidy(lm(Q3_CorrectNum ~ Party_Left, data = tmpdata))
broom::tidy(lm(Q3_CorrectNum ~ Party_Indp, data = tmpdata))

tmpdata %>%
  mutate(Party_Cat = case_when(Party_LDP == 1 ~ "LDP",
                               Party_Left == 1 ~ "Left",
                               Party_Indp == 1 ~ "Independent",
                               TRUE ~ "Others")) %>%
  select(Party_Cat, Q3_CorrectNum) %>%
  mutate(Knowledgable = if_else(Q3_CorrectNum >= 2, 1, 0)) %>%
  ungroup() %>%
  {{PartyKnowledge <<- .}} %>%
  group_by(Party_Cat) %>%
  summarize(Knowledgable_prop = mean(Knowledgable))

PartyKnowledge %>% group_by(Party_Cat) %>% summarize(Count = n())


## Create Table 1
res_control <- data.frame(matrix(0, nrow=length(use_order), ncol=3))
row.names(res_control) <- use_order
colnames(res_control) <- c("No/Lean", "Neutral", "Yes/Lean")

res_treated <- data.frame(matrix(0, nrow=length(use_order), ncol=3))
row.names(res_treated) <- use_order
colnames(res_treated) <- c("No/Lean", "Neutral", "Yes/Lean")

for(i in 1:length(use_order)){
  use_choices <- choice_rights[[use_order[i]]]$ordered
  
  change_choices_set <- c(
    Group=list(c("Group", "Group_new", "Group")),  # Treatment
    DV1=use_choices[1], DV2=use_choices[2]
  )
  
  tempdata <- data_change_choices(data, change_choices_set)
  
  data_control <- tempdata %>% filter(Group == 0 & DV1 != 0) %>%
    mutate(DV = case_when(
      DV1 == 4 ~ "No and Lean No",
      DV1 == 5 ~ "No and Lean No",
      DV1 == 3 ~ "Nutral",
      DV1 == 1 ~ "Yes and Lean Yes",
      DV1 == 2 ~ "Yes and Lean Yes"
    ))
  table_control <- table(data_control$DV)
  table_control <- round(table_control / sum(table_control)*100, 2)
  res_control[i, ] <- as.vector(table_control)
  
  data_treated <- tempdata %>% filter(Group == 1 & DV2 != 0)%>%
    mutate(DV = case_when(
      DV2 == 4 ~ "No and Lean No",
      DV2 == 5 ~ "No and Lean No",
      DV2 == 3 ~ "Nutral",
      DV2 == 1 ~ "Yes and Lean Yes",
      DV2 == 2 ~ "Yes and Lean Yes"
    ))
  table_treated <- table(data_treated$DV)
  table_treated <- round(table_treated / sum(table_treated)*100, 2)
  res_treated[i, ] <- as.vector(table_treated)
}

write.csv(res_control,
          paste0(savefolder_base,
                 "DescriptiveControl.csv"))

write.csv(res_treated,
          paste0(savefolder_base,
                 "DescriptiveTreated.csv"))

